//////////////////////////////////////////////////////////////////////////

#ifndef _BOOSTCONNECTIONIMPL_H_
#include "boostconnectionimpl.h"
#endif

//////////////////////////////////////////////////////////////////////////

#ifndef _CONNECTIONPACKET_H_
#include "connectionpacket.h"
#endif

#include <boost/bind.hpp>
#include <iostream>

//////////////////////////////////////////////////////////////////////////


// Alloc memory for static fields

// Buffer for UDP packets
char BoostConnectionImpl::m_RecvBufferUDP[MAX_BUFFER_SIZE];

// Working mode
BoostConnectionImpl::Mode BoostConnectionImpl::m_Mode = INVALID;

// IO_SERVICE ptr
io_service* BoostConnectionImpl::m_IOService = NULL;

// TCP acceptor ptr
tcp::acceptor* BoostConnectionImpl::m_Acceptor = NULL;

// UDP socket ptr
udp::socket* BoostConnectionImpl::m_UDPSocket = NULL;

// Flags
bool BoostConnectionImpl::m_Started = false;
bool BoostConnectionImpl::m_UDPAcceptStarted = false;

// Thread
boost::thread BoostConnectionImpl::m_Thread;

// UDP endpoint
udp::endpoint BoostConnectionImpl::m_UDPSendEndpoint;
udp::endpoint BoostConnectionImpl::m_UDPRecvEndpoint;

// Connections by endpoints
BoostConnectionImpl::ConnectionsByEndpoints
BoostConnectionImpl::m_ConnectionsByEndpoints;

// Endpoints by connections
BoostConnectionImpl::EndpointsByConnections
BoostConnectionImpl::m_EndpointsByConnections;



//////////////////////////////////////////////////////////////////////////


BoostConnectionImpl::BoostConnectionImpl()
// Initialize TCP socket
: m_TCPSocket( *GetIOService() )
{
}


//////////////////////////////////////////////////////////////////////////


BoostConnectionImpl::~BoostConnectionImpl()
{
}


//////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::SendTCP( const Packet& _packet )
{
    // Wrap packet to connection packet
    ConnectionPacket cp(_packet);

    // Pack connection packet code
    cp.PackCode();

    // Send data through TCP
    m_TCPSocket.async_send( boost::asio::buffer( cp.Data(), cp.Size() ),
        boost::bind( &BoostConnectionImpl::HandleSendTCP, this,
        boost::asio::placeholders::error,
        boost::asio::placeholders::bytes_transferred ) );

    // Run io_service thread
    Run();
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::SendUDP( const Packet& _packet )
{
    // Wrap packet to connection packet
    ConnectionPacket cp(_packet);

    // Pack connection packet code
    cp.PackCode();


    udp::endpoint sendEndpoint;

    if (m_Mode == SERVER)
    {
        sendEndpoint = GetClientEndpoint(this);
    }
    else if (m_Mode == CLIENT)
    {
        sendEndpoint = m_UDPSendEndpoint;
    }

    // Send data through UDP
    GetUDPSocket()->async_send_to( boost::asio::buffer( cp.Data(), cp.Size() ),
        endpoint,
        boost::bind( &BoostConnectionImpl::HandleSendUDP,
        this,
        boost::asio::placeholders::error,
        boost::asio::placeholders::bytes_transferred ) );

    // Run io_service thread
    Run();
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::StartReceiveTCP()
{
    // Recv data through TCP
    m_TCPSocket.async_receive( boost::asio::buffer(m_RecvBufferTCP, MAX_BUFFER_SIZE),
        boost::bind( &BoostConnectionImpl::HandleRecvTCP, this,
        boost::asio::placeholders::error,
        boost::asio::placeholders::bytes_transferred ) );

    Run();
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::StartReceiveUDP()
{
    // Recv data through UDP
    GetUDPSocket()->async_receive_from( boost::asio::buffer(m_RecvBufferUDP, MAX_BUFFER_SIZE),
        m_UDPRecvEndpoint,
        boost::bind( &BoostConnectionImpl::HandleRecvUDP,
        boost::asio::placeholders::error,
        boost::asio::placeholders::bytes_transferred ) );

    Run();
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::HandleSendTCP( const error_code& _error,
                                         size_t _bytesTransferred )
{
    if (!_error)
        ConnectionImpl::HandleSendTCP(_bytesTransferred);
    else
        std::cout << "ERROR SEND TCP: " << _error.message() << std::endl;
}


////////////////////////////////////////////////////////////////////////////////



void BoostConnectionImpl::HandleSendUDP( const error_code& _error, size_t _bytesTransferred )
{
    if (!_error)
        ConnectionImpl::HandleSendUDP(_bytesTransferred);
    else
        std::cout << "ERROR SEND UDP: " << _error.message() << std::endl;
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::HandleRecvTCP( const error_code& _error, size_t _bytesTransferred )
{
    if (!_error)
    {
        ConnectionPacket cp(m_RecvBufferTCP, _bytesTransferred);
        cp.UnPackCode();

        ConnectionImpl::HandleRecvTCP( cp.Data(), cp.Size() );
        StartReceiveTCP();
    }
    else
        std::cout << "ERROR RECV TCP: " << _error.message() << std::endl;
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::HandleRecvUDP( const error_code& _error, size_t _bytesTransferred )
{
    if (!_error)
    {
        ConnectionsByEndpoints::iterator it = m_ConnectionsByEndpoints.find(m_UDPRecvEndpoint);
        if ( it != m_ConnectionsByEndpoints.end() )
        {
            ConnectionPacket cp(m_RecvBufferUDP, _bytesTransferred);
            cp.UnPackCode();

            BoostConnectionImpl* conn = it->second;
            conn->ConnectionImpl::HandleRecvUDP( cp.Data(), cp.Size() );
        }
        StartReceiveUDP();
    }
    else if (_error.value() == 10022 || _error.value() == 10061)
        StartReceiveUDP();
    else
        std::cout << "ERROR RECV UDP: " << _error.message() << std::endl;
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::Connect( const char* _address )
{
    // Server can not connect to anybody (server is always server)
    if (m_Mode == SERVER)
        return;

    // Set client mode
    m_Mode = CLIENT;

    // Resolve address
    tcp::resolver resolver( *GetIOService() );
    tcp::resolver::query query(_address, m_PortStr);
    tcp::resolver::iterator endpoint_iterator = resolver.resolve(query);
    tcp::resolver::iterator end;

    error_code error = boost::asio::error::host_not_found;
    while (error && endpoint_iterator != end)
    {
        m_TCPSocket.close();
        m_TCPEndpoint = *endpoint_iterator;

        // Connect to server
        m_TCPSocket.connect(m_TCPEndpoint, error);
        ++endpoint_iterator;
    }
    if (!error)
    {
        // Also establish UDP connection
        udp::resolver resolver( *GetIOService() );
        udp::resolver::query query( udp::v4(), _address, m_PortStr );
        m_UDPRecvEndpoint = *resolver.resolve(query);
        m_UDPSendEndpoint = m_UDPRecvEndpoint;

        GetUDPSocket()->open( udp::v4() );
        GetUDPSocket()->send_to(boost::asio::buffer("connected", 9), m_UDPSendEndpoint);

        StartReceiveUDP();
        StartReceiveTCP();
    }
    else
    {
        std::cout << "ERROR CONNECT: " << error.message() << std::endl;
        throw boost::system::system_error(error);
    }
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::Accept()
{
    // Client can not accept connections from anybody (client is always client)
    if (m_Mode == CLIENT)
        return;

    // Set server mode
    m_Mode = SERVER;

    // Accept new TCP connections
    GetTCPAcceptor()->async_accept( m_TCPSocket,
        boost::bind( &BoostConnectionImpl::HandleAccept,
        this,
        boost::asio::placeholders::error ) );

    Run();
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::Close()
{
    m_TCPSocket.close();
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::HandleAccept( const error_code& _error )
{
    if (!_error)
    {
        udp::endpoint newConnEndpoint( m_TCPSocket.remote_endpoint().address(),
                                       m_TCPSocket.remote_endpoint().port() );

        m_ConnectionsByEndpoints[newConnEndpoint] = this;
        m_EndpointsByConnections[this] = newConnEndpoint;

        if (!m_UDPAcceptStarted)
        {
            // Also accept UDP packets
            GetUDPSocket()->open( udp::v4() );
            GetUDPSocket()->bind( udp::endpoint(udp::v4(), m_Port) );
            StartReceiveUDP();
            m_UDPAcceptStarted = true;
        }

        StartReceiveTCP();
        ConnectionImpl::HandleAccept();
    }
    else
        std::cout << "ERROR ACCEPT: " << _error.message() << std::endl;
}


////////////////////////////////////////////////////////////////////////////////


io_service* BoostConnectionImpl::GetIOService()
{
    if (!m_IOService)
        m_IOService = new io_service;

    return m_IOService;
}


////////////////////////////////////////////////////////////////////////////////


tcp::acceptor* BoostConnectionImpl::GetTCPAcceptor()
{
    if (!m_Acceptor)
        m_Acceptor = new tcp::acceptor( *GetIOService(), tcp::endpoint(tcp::v4(), m_Port) );

    return m_Acceptor;
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::Run()
{
    if (!m_Started)
        m_Thread = boost::thread(&BoostConnectionImpl::ThreadProc);
}


////////////////////////////////////////////////////////////////////////////////


void BoostConnectionImpl::ThreadProc()
{
    m_Started = true;
    GetIOService()->run();
    m_Started = false;
}


////////////////////////////////////////////////////////////////////////////////


udp::socket* BoostConnectionImpl::GetUDPSocket()
{
    if (!m_UDPSocket)
        m_UDPSocket = new udp::socket( *GetIOService() );

    return m_UDPSocket;
}


////////////////////////////////////////////////////////////////////////////////


udp::endpoint BoostConnectionImpl::GetUDPEndpoint( BoostConnectionImpl* _this )
{
    EndpointsByConnections::iterator it = m_EndpointsByConnections.find(_this);
    if ( it != m_EndpointsByConnections.end() )
        return it->second;
    else
        return udp::endpoint();
}


////////////////////////////////////////////////////////////////////////////////
